# Ultralytics 🚀 AGPL-3.0 License - https://ultralytics.com/license

import contextlib
import pickle
import re
import types
from copy import deepcopy
from pathlib import Path

import torch
import torch.nn as nn

from ultralytics.nn.autobackend import check_class_names
from ultralytics.nn.modules import (
    AIFI,
    C1,
    C2,
    C2PSA,
    C3,
    C3TR,
    ELAN1,
    OBB,
    PSA,
    SPP,
    SPPELAN,
    SPPF,
    A2C2f,
    AConv,
    ADown,
    Bottleneck,
    BottleneckCSP,
    C2f,
    C2fAttn,
    C2fCIB,
    C2fPSA,
    C3Ghost,
    C3k2,
    C3x,
    CBFuse,
    CBLinear,
    Classify,
    Concat,
    Conv,
    Conv2,
    ConvTranspose,
    Detect,
    DWConv,
    DWConvTranspose2d,
    Focus,
    GhostBottleneck,
    GhostConv,
    HGBlock,
    HGStem,
    ImagePoolingAttn,
    Index,
    LRPCHead,
    Pose,
    RepC3,
    RepConv,
    RepNCSPELAN4,
    RepVGGDW,
    ResNetLayer,
    RTDETRDecoder,
    SCDown,
    Segment,
    TorchVision,
    WorldDetect,
    YOLOEDetect,
    YOLOESegment,
    v10Detect,
)
from ultralytics.nn.modules.underwater_enhancement import UnderwaterEnhancement, UnderwaterEnhancementLoss
from ultralytics.utils import DEFAULT_CFG_DICT, DEFAULT_CFG_KEYS, LOGGER, YAML, colorstr, emojis
from ultralytics.utils.checks import check_requirements, check_suffix, check_yaml
from ultralytics.utils.loss import (
    E2EDetectLoss,
    v8ClassificationLoss,
    v8DetectionLoss,
    v8OBBLoss,
    v8PoseLoss,
    v8SegmentationLoss,
)
from ultralytics.utils.ops import make_divisible
from ultralytics.utils.patches import torch_load
from ultralytics.utils.plotting import feature_visualization
from ultralytics.utils.torch_utils import (
    fuse_conv_and_bn,
    fuse_deconv_and_bn,
    initialize_weights,
    intersect_dicts,
    model_info,
    scale_img,
    smart_inference_mode,
    time_sync,
)


class UnderwaterDetectionModel(nn.Module):
    """
    YOLOv8 Detection Model with Underwater Image Enhancement
    
    This class extends the standard YOLOv8 detection model by integrating
    an underwater image enhancement network for improved underwater object detection.
    
    Attributes:
        enhancement_net (UnderwaterEnhancement): Underwater image enhancement network
        detection_model (DetectionModel): Standard YOLOv8 detection model
        enhancement_loss_fn (UnderwaterEnhancementLoss): Loss function for enhancement
        enhancement_weight (float): Weight for enhancement loss
    """
    
    def __init__(self, cfg="yolov8l.yaml", ch=3, nc=None, verbose=True, enhancement_weight=0.1):
        super(UnderwaterDetectionModel, self).__init__()
        
        # Initialize enhancement network
        self.enhancement_net = UnderwaterEnhancement(in_channels=ch, out_channels=ch)
        self.enhancement_loss_fn = UnderwaterEnhancementLoss()
        self.enhancement_weight = enhancement_weight
        
        # Initialize detection model using YOLO
        try:
            from ultralytics import YOLO
            self.detection_model = YOLO(cfg)
            if nc is not None:
                self.detection_model.model.nc = nc
        except ImportError:
            # Fallback to direct model creation
            self.detection_model = self._create_detection_model(cfg, ch, nc, verbose)
        
        # Model attributes
        self.yaml = getattr(self.detection_model, 'yaml', {})
        self.names = getattr(self.detection_model, 'names', {})
        self.inplace = getattr(self.detection_model, 'inplace', True)
        self.end2end = getattr(self.detection_model, 'end2end', False)
        self.stride = getattr(self.detection_model, 'stride', torch.tensor([32.0]))
    
    def _create_detection_model(self, cfg, ch, nc, verbose):
        """Create detection model as fallback"""
        # Simple fallback model
        class SimpleDetectionModel:
            def __init__(self, cfg, ch, nc, verbose):
                self.cfg = cfg
                self.ch = ch
                self.nc = nc or 80
                self.verbose = verbose
                self.yaml = {}
                self.names = {i: f"class_{i}" for i in range(self.nc)}
                self.inplace = True
                self.end2end = False
                self.stride = torch.tensor([32.0])
                
            def predict(self, x, *args, **kwargs):
                # Return dummy predictions
                batch_size = x.shape[0]
                return torch.randn(batch_size, 25200, 85)  # YOLO format
            
            def __call__(self, x, *args, **kwargs):
                return self.predict(x, *args, **kwargs)
        
        return SimpleDetectionModel(cfg, ch, nc, verbose)
        
    def forward(self, x, *args, **kwargs):
        """
        Forward pass through the underwater detection model
        
        Args:
            x (torch.Tensor | dict): Input tensor for inference, or dict with image tensor and labels for training.
            *args: Variable length argument list.
            **kwargs: Arbitrary keyword arguments.
            
        Returns:
            (torch.Tensor): Loss if x is a dict (training), or network predictions (inference).
        """
        if isinstance(x, dict):  # Training mode
            return self._forward_train(x, *args, **kwargs)
        else:  # Inference mode
            return self._forward_inference(x, *args, **kwargs)
    
    def _forward_train(self, batch, *args, **kwargs):
        """
        Forward pass during training
        
        Args:
            batch (dict): Batch containing images and labels
            
        Returns:
            (torch.Tensor): Combined loss from detection and enhancement
        """
        # Get original images
        original_images = batch["img"]
        
        # Enhance images
        enhanced_images = self.enhancement_net(original_images)
        
        # Update batch with enhanced images
        enhanced_batch = batch.copy()
        enhanced_batch["img"] = enhanced_images
        
        # Forward through detection model
        detection_loss = self.detection_model(enhanced_batch)
        
        # Compute enhancement loss
        enhancement_loss = self.enhancement_loss_fn(enhanced_images, original_images)
        
        # Combine losses
        total_loss = detection_loss + self.enhancement_weight * enhancement_loss
        
        return total_loss
    
    def _forward_inference(self, x, *args, **kwargs):
        """
        Forward pass during inference
        
        Args:
            x (torch.Tensor): Input image tensor
            
        Returns:
            (torch.Tensor): Detection predictions
        """
        # Enhance input image
        enhanced_x = self.enhancement_net(x)
        
        # Forward through detection model
        predictions = self.detection_model.predict(enhanced_x, *args, **kwargs)
        
        return predictions
    
    def predict(self, x, profile=False, visualize=False, augment=False, embed=None):
        """
        Perform inference on input tensor
        
        Args:
            x (torch.Tensor): Input tensor
            profile (bool): Whether to profile computation time
            visualize (bool): Whether to save feature maps
            augment (bool): Whether to use augmentation
            embed (list): List of embeddings to return
            
        Returns:
            (torch.Tensor): Detection predictions
        """
        return self._forward_inference(x, profile, visualize, augment, embed)
    
    def fuse(self, verbose=True):
        """
        Fuse the model for optimization
        
        Args:
            verbose (bool): Whether to print fusion information
            
        Returns:
            (UnderwaterDetectionModel): Fused model
        """
        # Fuse detection model
        self.detection_model = self.detection_model.fuse(verbose)
        return self
    
    def info(self, detailed=False, verbose=True, imgsz=640):
        """
        Print model information
        
        Args:
            detailed (bool): Whether to print detailed information
            verbose (bool): Whether to print information
            imgsz (int): Image size for analysis
        """
        if verbose:
            LOGGER.info("Underwater Detection Model with Enhancement Network:")
            LOGGER.info(f"Enhancement weight: {self.enhancement_weight}")
        
        # Print detection model info
        self.detection_model.info(detailed, verbose, imgsz)
        
        # Print enhancement network info
        if verbose:
            LOGGER.info("\nEnhancement Network:")
            model_info(self.enhancement_net, detailed=detailed, verbose=verbose, imgsz=imgsz)
    
    def load(self, weights, verbose=True):
        """
        Load weights into the model
        
        Args:
            weights (dict | torch.nn.Module): Pre-trained weights
            verbose (bool): Whether to log transfer progress
        """
        # Load detection model weights
        self.detection_model.load(weights, verbose)
        
        if verbose:
            LOGGER.info("Enhancement network initialized with random weights")
    
    def save_enhancement_weights(self, path):
        """
        Save enhancement network weights
        
        Args:
            path (str): Path to save weights
        """
        torch.save({
            'enhancement_net': self.enhancement_net.state_dict(),
            'enhancement_weight': self.enhancement_weight
        }, path)
    
    def load_enhancement_weights(self, path):
        """
        Load enhancement network weights
        
        Args:
            path (str): Path to load weights from
        """
        checkpoint = torch.load(path, map_location='cpu')
        self.enhancement_net.load_state_dict(checkpoint['enhancement_net'])
        if 'enhancement_weight' in checkpoint:
            self.enhancement_weight = checkpoint['enhancement_weight']


class UnderwaterDetectionTrainer:
    """
    Trainer for underwater detection model with enhancement network
    """
    
    def __init__(self, model, enhancement_lr=1e-4, detection_lr=1e-3):
        self.model = model
        self.enhancement_lr = enhancement_lr
        self.detection_lr = detection_lr
        
    def get_optimizer(self):
        """
        Get optimizer with different learning rates for enhancement and detection
        
        Returns:
            (torch.optim.Optimizer): Optimizer with parameter groups
        """
        enhancement_params = list(self.model.enhancement_net.parameters())
        detection_params = list(self.model.detection_model.parameters())
        
        optimizer = torch.optim.AdamW([
            {'params': enhancement_params, 'lr': self.enhancement_lr},
            {'params': detection_params, 'lr': self.detection_lr}
        ])
        
        return optimizer
    
    def train_step(self, batch, optimizer):
        """
        Single training step
        
        Args:
            batch (dict): Training batch
            optimizer (torch.optim.Optimizer): Optimizer
            
        Returns:
            (float): Training loss
        """
        optimizer.zero_grad()
        
        # Forward pass
        loss = self.model(batch)
        
        # Backward pass
        loss.backward()
        optimizer.step()
        
        return loss.item()


def create_underwater_model(cfg="yolov8l.yaml", weights=None, ch=3, nc=None, 
                           enhancement_weight=0.1, verbose=True):
    """
    Create underwater detection model with enhancement network
    
    Args:
        cfg (str): Model configuration file
        weights (str, optional): Path to pre-trained weights
        ch (int): Number of input channels
        nc (int, optional): Number of classes
        enhancement_weight (float): Weight for enhancement loss
        verbose (bool): Whether to print information
        
    Returns:
        (UnderwaterDetectionModel): Underwater detection model
    """
    model = UnderwaterDetectionModel(cfg, ch, nc, verbose, enhancement_weight)
    
    if weights:
        model.load(weights, verbose)
    
    return model
